import time, hmac, json, base64, hashlib

# 设置headers 加密算法的配置
headers = {
    "alg": "HS256",
    "typ": "JWT",
}
# 密钥
salt = '(((@*###$$$%%%*@)))'


def create_token(username):
    # 设置过期时间为6小时 60 * 60 * 6
    exp = round(time.time()) + 60 * 60 * 6

    payload = {
        "exp": exp,
        "username": username
    }

    first = base64.b64encode(json.dumps(headers).encode('utf-8')).decode('utf-8')
    second = base64.b64encode(json.dumps(payload).encode('utf-8')).decode('utf-8')
    # 第一第二部分合并
    first_second = f"{first}.{second}"
    # 生成签名
    third = base64.b64encode(
        hmac.new(salt.encode('utf-8'), first_second.encode('utf-8'), hashlib.sha256).digest()).decode('utf-8')
    # token
    token = ".".join([first, second, third])
    return token


def check_token(token: str):
    if not token.count('.'):
        return None
    token_list = token.split('.')
    h_str = token_list[0]
    p_str = token_list[1]
    s_str = token_list[2]
    payload = json.loads(base64.b64decode(p_str))
    # 第三部分
    third = base64.b64encode(
        hmac.new(salt.encode('utf-8'), (h_str + '.' + p_str).encode('utf-8'), hashlib.sha256).digest()).decode('utf-8')
    if s_str == third:
        return payload
    else:
        return None
